﻿#pragma once

#include <iostream>
#include <cassert>
#include <cstdlib>
using namespace std;

namespace Tree
{
	enum Color
	{
		BLACK,
		RED
	};

	template<class K>
	struct Node
	{
		Node(K key)
			:_val(key)
			, _col (RED)
			,_left(nullptr)
			,_right(nullptr)
			,_parent(nullptr)
		{}

		int _Height(Node<K>* root)
		{
			if (root == nullptr)
				return 0;
			int leftHeight = _Height(root->_left);
			int rightHeight = _Height(root->_right);
			return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
		}

		void _InOrder(Node<K>* _root)
		{
			if (_root == nullptr)
			{
				return;
			}

			_InOrder(_root->_left);
			cout << _root->_val << endl;
			_InOrder(_root->_right);
		}

		bool check_tree(Node<K>* root, int count, int height)
		{
			if (root == nullptr)
			{
				if (count == height)
				{
					return true;
				}
				else
				{
					return false;
				}
			}

			if (root->_col == RED)
			{
				Node<K>* left = root->_left;
				Node<K>* right = root->_right;

				if (left && left->_col == RED)
				{
					return false;
				}
				if (right && right->_col == RED)
				{
					return false;
				}
			}
			else
			{
				++count;
			}

			return check_tree(root->_left, count, height) && check_tree(root->_right, count, height);
		}

		~Node()
		{}

		K _val;
		Color _col;
		Node<K>* _left;
		Node<K>* _right;
		Node<K>* _parent;
	};

	template<class K>
	class Red_Black_Tree
	{
	public:
		typedef Node<K> Node;

		int Height()
		{
			return _root->_Height(_root);
		}

		void Inorder()
		{
			_root->_InOrder(_root);
		}

		bool check()
		{
			if (_root->_col == BLACK)
				return false;

			Node* cur = _root;
			int height = 0;
			while (cur)
			{
				if (cur->_col == BLACK)
					++height;
				cur = cur->_left;
			}

			return (_root->check_tree(_root, 0, height));
		}

		// 左单旋
		void RotateL(Node* parent)
		{
			Node* SubR = parent->_right;
			Node* SubRL = SubR->_left;
			Node* pParent = parent->_parent;

			parent->_right = SubRL;
			if (SubRL)
				SubRL->_parent = parent;

			SubR->_left = parent;

			SubR->_parent = pParent;
			if (pParent == nullptr)
			{
				_root = SubR;
			}
			else
			{
				if (pParent->_left == parent)
				{
					pParent->_left = SubR;
				}
				else
				{
					pParent->_right = SubR;
				}
			}
		}
		
		// 右单旋
		void RotateR(Node* parent)
		{
			Node* SubL = parent->_left;
			Node* SubLR = SubL->_right;
			Node* pParent = parent->_parent;

			parent->_left = SubLR;
			if (SubLR)
				SubLR->_parent = parent;

			SubL->_left = parent;
			parent->_parent = SubL;

			SubL->_parent = pParent;
			if (pParent == nullptr)
			{
				_root = SubL;
			}
			else
			{
				if (pParent->_left == parent)
				{
					pParent->_left = SubL;
				}
				else
				{
					pParent->_right = SubL;
				}
			}
		}

		// 旋转代码的实现跟AVL树是⼀样的，只是不需要更新平衡因⼦ 

		bool answer_Insert(const K& kv)
		{
			if (_root == nullptr)
			{
				_root = new Node(kv);
				_root->_col = BLACK;
				return true;
			}
			Node* parent = nullptr;
			Node* cur = _root;
			while (cur)
			{
				if (cur->_val < kv)
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (cur->_val > kv)
				{
					parent = cur;
					cur = cur->_left;
				}
				else

				{
					return false;
				}
			}
			cur = new Node(kv);
			// 新增结点。颜⾊红⾊给红⾊ 
			cur->_col = RED;
			if (parent->_val < kv)
			{
				parent->_right = cur;
			}
			else
			{
				parent->_left = cur;
			}
			cur->_parent = parent;
			while (parent && parent->_col == RED)
			{
				Node* grandfather = parent->_parent;
				if (grandfather == nullptr)
					break;

				// g

				// p u

				if (parent == grandfather->_left)
				{
					Node* uncle = grandfather->_right;
					if (uncle && uncle->_col == RED) {
						// u存在且为红 -》变⾊再继续往上处理 
						parent->_col = uncle->_col = BLACK;
						grandfather->_col = RED;
						cur = grandfather;
						parent = cur->_parent;
					}
					else

					{
						// u存在且为⿊或不存在 -》旋转+变⾊ 
						if (cur == parent->_left)
						{
							// g

							// p u

							//c

							//单旋 
							RotateR(grandfather);
							parent->_col = BLACK;
							grandfather->_col = RED;
						}
						else

						{
							// g

							// p u

							// c

							//双旋 
							RotateL(parent);
							RotateR(grandfather);
							cur->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}
				}
				else

				{
					// g

					// u p

					Node* uncle = grandfather->_left;
					// 叔叔存在且为红，-》变⾊即可 
					if (uncle && uncle->_col == RED)
					{
						parent->_col = uncle->_col = BLACK;
						grandfather->_col = RED;
						// 继续往上处理 
						cur = grandfather;
						parent = cur->_parent;
					}
					else // 叔叔不存在，或者存在且为⿊ 
					{
						// 情况⼆：叔叔不存在或者存在且为⿊ 
						// 旋转+变⾊ 
						// g

						// u p

						// c

						if (cur == parent->_right)
						{
							RotateL(grandfather);
							parent->_col = BLACK;
							grandfather->_col = RED;
						}
						else

						{
							// g

							// u p

							// c

							RotateR(parent);
							RotateL(grandfather);
							cur->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}
				}
			}
			_root->_col = BLACK;
			return true;
		}

		bool Insert(const K& key)
		{
			if (_root == nullptr)
			{
				_root = new Node(key);
				return true;
			}

			Node* cur = _root;
			Node* parent = nullptr;
			while (cur)
			{
				if (key > cur->_val)
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (key < cur->_val)
				{
					parent = cur;
					cur = cur->_left;
				}
				else
				{
					return false;
				}
			}

			cur = new Node(key);
			if (key > parent->_val)
				parent->_right = cur;
			else
				parent->_left = cur;

			cur->_parent = parent;
			cur->_col = RED;

			while (parent && parent->_col == RED)
			{
				Node* grandfather = parent->_parent;
				Node* Uncle = nullptr;
				if (grandfather->_left == parent)
				{
					Uncle = grandfather->_right;
					if(Uncle && Uncle->_col == RED)
					{
						grandfather->_col = RED;
						Uncle->_col = parent->_col = BLACK;
					}
					else
					{
						if (cur == parent->_left)
						{
							RotateR(grandfather);
							grandfather->_col = RED;
							parent->_col = BLACK;
						}
						else
						{
							RotateL(parent);
							RotateR(grandfather);
							cur->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}
				}
				else
				{
					Uncle = grandfather->_left;
					if(Uncle && Uncle->_col == RED)
					{
						grandfather->_col = RED;
						Uncle->_col = parent->_col = BLACK;
					}
					else
					{
						if (cur == parent->_left)
						{
							RotateL(grandfather);
							grandfather->_col = RED;
							parent->_col = BLACK;
						}
						else
						{
							RotateR(parent);
							RotateL(grandfather);
							cur->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}
				}
				cur = grandfather;
				parent = cur->_parent;
			}
			_root->_col = BLACK;
			return true;
		}
	private:
		Node* _root = nullptr;
	};

	void test()
	{
		Red_Black_Tree<int> t;
		srand(time(NULL));

		int n = 1000;

		clock_t start = clock();
		while (n--)
		{
			int i = rand() + n;
			t.answer_Insert(i);
		}
		clock_t end = clock();
		// cout << end - start << endl;
		cout << t.Height() << endl;
		//t.Inorder();
		cout << t.check() << endl;
	}
}
